# 安装依赖 pip install langchain_community
import os

from langchain_openai import AzureChatOpenAI
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, ConfigurableFieldSpec

os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "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"

os.environ["AZURE_OPENAI_ENDPOINT"] = "http://menshen.test.xdf.cn"
os.environ["OPENAI_API_KEY"] = "c8575027653b42b1b47747f0b4ab135b"
os.environ["OPENAI_API_TYPE"] = "azure"
os.environ["OPENAI_API_VERSION"] = "2023-05-15"

llm = AzureChatOpenAI(
    deployment_name="gpt-4o",
    model_name="gpt-4o",
    temperature=0
)

prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个{role}专家，请用你的专业知识回答问题"),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{question}")
])

#  prompt模板链式调用
chain = prompt | llm

# 存储聊天记录
store = {}


def get_session_history(user_id: str, conversation_id: str):
    if (user_id, conversation_id) not in store:
        store[(user_id, conversation_id)] = ChatMessageHistory()
    return store[(user_id, conversation_id)]


do_message = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key="question",  # 每次聊天时发送消息的key
    history_messages_key="chat_history",
    history_factory_config=[
        ConfigurableFieldSpec(
            id="user_id",
            annotation=str,
            name="用户id",
            default="",
            description="用户id",
            is_shared=True
        ),
        ConfigurableFieldSpec(
            id="conversation_id",
            annotation=str,
            name="会话id",
            default="",
            description="会话id",
            is_shared=True
        )
    ]

)

config = {'configurable': {'user_id': '123', 'conversation_id': 'zs123'}}

# 第1轮
resp = do_message.invoke(
    {
        "role": "互联网技术研发",
        "question": "学习langchain都有哪些知识点"
    },
    config
)
print(resp.content)

# 第2轮
resp = do_message.invoke(
    {
        "role": "互联网技术研发",
        "question": "什么"
    },
    config
)
print(resp.content)

# 第3轮 流式返回
for resp in do_message.stream(
        {
            "role": "互联网技术研发",
            "question": "再说一遍"
        },
        config
):
    print(resp.content, end='-')
